from random import sample, random, choice, randint
from math import ceil, log
import pandas as pd
import numpy as np
import time
import rrcf
import matplotlib as mpl
mpl.use('TkAgg')
from sklearn.ensemble import IsolationForest
from matplotlib import pyplot as plt
# Read data
OilWell = pd.read_csv('tp152.csv', index_col=0)
OilWell.index = pd.to_datetime(OilWell.index)
data = OilWell['pipe'].astype(float).values

# Create events
events = {
'event #1'          : ('2020-1-03 16:25:00',
                      '2020-1-03 16:30:00'),
'event #2'          : ('2020-1-04 12:28:00',
                      '2020-1-04 12:38:00'),
'event #3'          : ('2020-1-04 23:28:00',
                      '2020-1-04 23:33:00'),
'event #4'          : ('2020-1-05 05:24:00',
                      '2020-1-05 05:30:00'),
'event #5'          : ('2020-1-05 06:53:00',
                      '2020-1-05 06:58:00'),
'event #6'          : ('2020-1-05 09:28:00',
                      '2020-1-05 09:33:00'),
'event #7'          : ('2020-1-05 11:23:00',
                      '2020-1-05 11:28:00'),
'event #8'          : ('2020-1-06 08:55:00',
                      '2020-1-06 09:00:00'),
'event #9'          : ('2020-1-07 00:36:00',
                      '2020-1-07 00:41:00')
}
OilWell['event'] = np.zeros(len(OilWell))
for event, duration in events.items():
    start, end = duration
    OilWell.loc[start:end, 'event'] = 1

# Next, we run the RRCF algorithm and compute the CoDisp for each point.

# Set tree parameters
num_trees = 200
shingle_size = 30
tree_size = 500

# Use the "shingle" generator to create rolling window
points = rrcf.shingle(data, size=shingle_size)
points = np.vstack([point for point in points])
n = points.shape[0]
sample_size_range = (n // tree_size, tree_size)

time_start=time.time()

#from utils import run_time
class Node(object):
    def __init__(self, size):
        """Node class to build tree leaves
        Keyword Arguments:
            size {int} -- Node size (default: {None})
        """
        # Node size
        self.size = size
        # Feature to split
        self.split_feature = None
        # Split point
        self.split_point = None
        # Left child node
        self.left = None
        # Right child node
        self.right = None
class IsolationTree(object):
    def __init__(self, X, n_samples, max_depth):
        """Isolation Tree class
        Arguments:
            X {list} -- 2d list with int or float
            n_samples {int} -- Subsample size
            max_depth {int} -- Maximum height of isolation tree
        """
        self.height = 0
        # In case of n_samples is greater than n
        n = len(X)
        if n_samples > n:
            n_samples = n
        # Root node
        self.root = Node(n_samples)
        # Build isolation tree
        self._build_tree(X, n_samples, max_depth)
    def _get_split(self, X, idx, split_feature):
        """Randomly choose a split point
        Arguments:
            X {list} -- 2d list object with int or float
            idx {list} -- 1d list object with int
            split_feature {int} -- Column index of X
        Returns:
            int -- split point
        """
        # The split point should be greater than min(X[feature])
        
        unique = set(map(lambda i: X[i][split_feature], idx))
#        print("unique:",unique)
#        for iii in range(10):
#            print("idx,split_feature,X[",iii,"][split_feature]:",idx,split_feature,X[iii][split_feature])
        # Cannot split
        if len(unique) == 1:
            return None
        unique.remove(min(unique))
        x_min, x_max = min(unique), max(unique)
        # Caution: random() -> x in the interval [0, 1).
#        print("打印随机值:",random() * (x_max - x_min) + x_min)
#        print("----------------------------------------------")
        return random() * (x_max - x_min) + x_min
    def _build_tree(self, X, n_samples, max_depth):
        """The current node data space is divided into 2 sub space: less than the
        split point in the specified dimension on the left child of the current node,
        put greater than or equal to split point data on the current node's right child.
        Recursively construct new child nodes until the data cannot be splitted in the
        child nodes or the child nodes have reached the max_depth.
        Arguments:
            X {list} -- 2d list object with int or float
            n_samples {int} -- Subsample size
            max_depth {int} -- Maximum depth of IsolationTree
        """
        # Dataset shape
        m = len(X[0])
        n = len(X)
        # Randomly selected sample points into the root node of the tree
        idx = sample(range(n), n_samples)
        # Depth, Node and idx
        que = [[0, self.root, idx]]
        # BFS
        while que and que[0][0] <= max_depth:
            depth, nd, idx = que.pop(0)
            # Stop split if X cannot be splitted
            nd.split_feature = choice(range(m))
            nd.split_point = self._get_split(X, idx, nd.split_feature)
            if nd.split_point is None:
                continue
            # Split
            idx_left = []
            idx_right = []
            while idx:
                i = idx.pop()
                xi = X[i][nd.split_feature]
                if xi < nd.split_point:
                    idx_left.append(i)
                else:
                    idx_right.append(i)
            # Generate left and right child
            nd.left = Node(len(idx_left))
            nd.right = Node(len(idx_right))
            # Put the left and child into the que and depth plus one
            que.append([depth+1, nd.left, idx_left])
            que.append([depth+1, nd.right, idx_right])
        # Update the height of IsolationTree
        self.height = depth
    def _predict(self, xi):
        """Auxiliary function of predict.
        Arguments:
            xi {list} -- 1D list with int or float
        Returns:
            int -- the depth of the node which the xi belongs to
        """
        # Search xi from the IsolationTree until xi is at an leafnode
        nd = self.root
        count=0
        depth = 0
        while nd.left and nd.right:
            if xi[nd.split_feature] < nd.split_point:
                nd = nd.left                
            else:
                nd = nd.right
            #打印树结构和树高度等信息
#            print("nd.split_feature,xi[nd.split_feature],nd.split_point:",nd.split_feature,xi[nd.split_feature],nd.split_point)
            depth += 1
#        print("depth,nd.size:",depth,nd.size)
        return depth, nd.size
    
class IsolationForest2(object):
    def __init__(self):
        """IsolationForest, randomly build some IsolationTree instance,
        and the average score of each IsolationTree
        Attributes:
        trees {list} -- 1d list with IsolationTree objects
        ajustment {float}
        """
        self.trees = None
        self.adjustment = None  # TBC
    def fit(self, X, n_samples=50, max_depth=10, n_trees=50):
        """Build IsolationForest with dataset X
        Arguments:
            X {list} -- 2d list with int or float
        Keyword Arguments:
            n_samples {int} -- According to paper, set number of samples to 256 (default: {256})
            max_depth {int} -- Tree height limit (default: {10})
            n_trees {int} --  According to paper, set number of trees to 100 (default: {100})
        """
        self.adjustment = self._get_adjustment(n_samples)
        self.trees = [IsolationTree(X, n_samples, max_depth)
                      for _ in range(n_trees)]
    def _get_adjustment(self, node_size):
        """Calculate adjustment according to the formula in the paper.
        Arguments:
            node_size {int} -- Number of leaf nodes
        Returns:
            float -- ajustment
        """
        if node_size > 2:
            i = node_size - 1
            ret = 2 * (log(i) + 0.5772156649) - 2 * i / node_size
        elif node_size == 2:
            ret = 1
        else:
            ret = 0
#        print("ret: ",ret)
        return ret

    def _predict(self, xi):
        """Auxiliary function of predict.
        Arguments:
            xi {list} -- 1d list object with int or float
        Returns:
            list -- 1d list object with float
        """
        # Calculate average score of xi at each tree
        score = 0
        n_trees = len(self.trees)
        for tree in self.trees:
            depth, node_size = tree._predict(xi) #调用孤立树的方法
#            print("xi,depth:",xi,xi[0],depth)
            score += (depth + self._get_adjustment(node_size))
        score = score / n_trees
        # Scale
#       print("n_trees ,self.adjustment,score ,2 ** -(score / self.adjustment):",n_trees ,self.adjustment,score,2 ** -(score / self.adjustment))
        return 2 ** -(score / self.adjustment)
    def predict(self, X):
        """Get the prediction of y.
        Arguments:
            X {list} -- 2d list object with int or float
        Returns:
            list -- 1d list object with float
        """
        return [self._predict(xi) for xi in X]
#@run_time


def main():
    #
    clf = IsolationForest2()
    clf.fit(points, n_samples=100, n_trees=200)

    ppp=clf.predict(points)
    ppp=np.array(ppp)
    contamination = OilWell['event'].sum()/len(OilWell)
    IF = IsolationForest(n_estimators=num_trees,
                         contamination=contamination,
                         random_state = 0)
    IF.fit(points)
    if_scores = IF.score_samples(points)


    ppp=pd.Series(ppp,index=(OilWell.iloc[(shingle_size - 1):].index))

    if_scores = pd.Series(-if_scores,
                          index=(OilWell
                                 .iloc[(shingle_size - 1):]
                                 .index))

    ppp = ((ppp - ppp.min())
                  / (ppp.max() - ppp.min()))
    if_scores = ((if_scores - if_scores.min())
                  / (if_scores.max() - if_scores.min()))
#
    np.savetxt('scores_result.txt ', np.hstack([ppp, if_scores]))

    fig, ax = plt.subplots(2, figsize=(15, 6.0))
    (OilWell['pipe'] / 1.0).plot(ax=ax[0], color='0.2', alpha=0.9)
    if_scores.plot(ax=ax[1], color='#7EBDE6', alpha=0.95, label='IsolationForest scores')
    ppp.plot(ax=ax[1], color='#E8685D', alpha=0.95, label='Adjusted IsolationForest scores')

    ax[1].legend(frameon=True, loc=2, fontsize=12)

    for event, duration in events.items():
        start, end = duration
        ax[0].axvspan(start, end, alpha=0.5, color='springgreen')

    ax[0].set_xlabel('')
    ax[1].set_xlabel('')

    ax[0].set_ylabel('Pipeline pressure(Mpa)', size=10)
    ax[1].set_ylabel('Normalized anomaly score', size=10)
    ax[0].grid(alpha=0.2)
    ax[0].set_title('Oil-well pressure monitoring', size=12)
    ax[0].xaxis.set_ticklabels([])

    ax[0].set_xlim(OilWell.index[0], OilWell.index[-1])
    ax[1].set_xlim(OilWell.index[0], OilWell.index[-1])
    plt.grid(alpha=0.2)
    plt.axhline(y=0.32, linestyle='--', linewidth = 1.0, color='blue', alpha=0.9)
    plt.axhline(y =0.45, linestyle='-.',linewidth = 1.0, color='limegreen', alpha=0.9)

    plt.tight_layout()

    plt.show()
#


if __name__ == "__main__":
    main()